Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops #648

xumingkuan · 2020-03-23T23:59:49Z

Related issue id = #631

Also make mgpcg_advanced.py able to run with self.dim = 2.

xumingkuan · 2020-03-24T01:42:33Z

It gives me TypeError: list indices must be integers or slices, not Expr whenever I write sth like __ndrange.acc_dimensions[__grouped_I + 1]. I think it's necessary to have a variable in some indices in the template string in transformer.py. How can I walk around it?

xumingkuan · 2020-03-24T01:44:01Z

python/taichi/lang/transformer.py

+            loop_body = t_loop.body
+            inner_loop_body = loop_body[1].body
+            inner_loop_body[1].body[0].value = self.parse_expr(
+                '__I // __ndrange.acc_dimensions[__grouped_I + 1]')


I tried using self.parse_expr to avoid the Expr index in the template string but it didn't work.

yuanming-hu · 2020-03-24T02:01:37Z

python/taichi/lang/transformer.py

+            loop_body += node.body
+
+            node = ast.copy_location(t, node)
+            return self.visit(node)  # further translate as a range for


Here self.visit will further translate __grouped_I_tmp = 0 into __grouped_I_tmp = ti.expr_init(0) and you get an expr. The solution is to first visit the node body, and then do your own transform so that what you write is what you end up with.

After visiting the node body (translating as a range for), how can I locate the statements I want to modify?

An easier way: see the trick here and protect your statements with ti.static(1) so that it won't get further translated

taichi/python/taichi/lang/transformer.py

Line 340 in 063542f

if ti.static(1):

Sorry, to be clear: try __grouped_I_tmp = ti.static(0) so that this statement won't get transformed into __grouped_I_tmp = ti.expr_init(0).

After visiting the node body (translating as a range for), how can I locate the statements I want to modify?

In my understanding, the only transform you have to do here is to prepend some statements to the loop body so that you can get something like I. You don't need to modify any existing statements in the old loop body.

yuanming-hu · 2020-03-24T13:15:59Z

An easy way to debug this is to enable print_processed=True in ti.init or @all_archs_with, so that you can look at the transformed AST and figure out why you get TypeError: list indices must be integers or slices, not Expr :-)

xumingkuan · 2020-03-24T17:59:29Z

I need __ndrange.acc_dimensions[__grouped_I + 1] and __ndrange.bounds[__grouped_I][0].

However, __ndrange.acc_dimensions is a list, while __grouped_I + 1 is an Expr.

As the error says, list indices are not allowed to be Expr.

Maybe I should make use of GroupedNDRange to obtain a Vector but I don't know how to use it. Shall we have a Skype chat now?

yuanming-hu · 2020-03-24T18:13:48Z

Shall we have a Skype chat now?

Sure.

xumingkuan · 2020-03-24T18:16:39Z

After preprocessing:
def test():
  if ti.static(1):
    __ndrange = ti.expr_init(ti.ndrange((x0, y0), (x1, y1)))
    I = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions)))
    if 1:
      __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
      ___begin = ti.Expr(0)
      ___end = ti.Expr(ti.subscript(__ndrange.acc_dimensions, 0))
      ___begin = ti.cast(___begin, ti.i32)
      ___end = ti.cast(___end, ti.i32)
      ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
          ___end.ptr)
      __I = ti.expr_init(__ndrange_I)
      if 1:
        __grouped_I = ti.Expr(ti.core.make_id_expr(''))
        ___begin = ti.Expr(0)
        ___end = ti.Expr(len(__ndrange.dimensions))
        ___begin = ti.cast(___begin, ti.i32)
        ___end = ti.cast(___end, ti.i32)
        ti.core.begin_frontend_range_for(__grouped_I.ptr, ___begin.ptr,
            ___end.ptr)
        __grouped_I_tmp = ti.static(0)
        if 1:
          __cond = ti.chain_compare([__grouped_I + 1, len(__ndrange.
              dimensions)], ['Lt'])
          ti.core.begin_frontend_if(ti.Expr(__cond).ptr)
          ti.core.begin_frontend_if_true()
          __grouped_I_tmp = ti.expr_init(233)
          del __grouped_I_tmp
          ti.core.pop_scope()
          ti.core.begin_frontend_if_false()
          __grouped_I_tmp = ti.expr_init(__I)
          del __grouped_I_tmp
          ti.core.pop_scope()
        __tmp = ti.expr_init(233)
        if 1:
          __cond = ti.chain_compare([__grouped_I + 1, len(__ndrange.
              dimensions)], ['Lt'])
          ti.core.begin_frontend_if(ti.Expr(__cond).ptr)
          ti.core.begin_frontend_if_true()
          __I.assign(233)
          ti.core.pop_scope()
          ti.core.begin_frontend_if_false()
          ti.core.pop_scope()
        del __tmp
        ti.core.end_frontend_range_for()
        del __grouped_I
      ti.subscript(val, I).assign(ti.subscript(I, 0) + ti.subscript(I, 1) * 2)
      del __I
      ti.core.end_frontend_range_for()
      del __ndrange_I
    del I
    del __ndrange

xumingkuan · 2020-03-24T18:37:46Z

After preprocessing:
def test():
  if ti.static(1):
    __ndrange = ti.ndrange((x0, y0), (x1, y1))
    I = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions)))
    ___begin = ti.Expr(0)
    ___end = __ndrange.acc_dimensions[0]
    ___begin = ti.cast(___begin, ti.i32)
    ___end = ti.cast(___end, ti.i32)
    __ndrange_I = ti.Expr(0)
    ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr, ___end.ptr)
    __I = __ndrange_I
    for __grouped_I in range(len(__ndrange.dimensions)):
      __grouped_I_tmp = 0
      if __grouped_I + 1 < len(__ndrange.dimensions):
        __grouped_I_tmp = __I // __ndrange.acc_dimensions[__grouped_I + 1]
      else:
        __grouped_I_tmp = __I
      ti.subscript(I, __grouped_I).assign(__grouped_I_tmp + __ndrange.
          bounds[__grouped_I][0])
      if __grouped_I + 1 < len(__ndrange.dimensions):
        __I = __I - __grouped_I_tmp * __ndrange.acc_dimensions[__grouped_I + 1]
    ti.subscript(val, I).assign(ti.subscript(I, 0) + ti.subscript(I, 1) * 2)
    ti.core.end_frontend_range_for()

# Conflicts: # python/taichi/lang/transformer.py

…dim=2

xumingkuan · 2020-03-24T20:53:58Z

Ready for review.

yuanming-hu

Great job! Now visit_For is much easier to read. Note that currently visit_For is still a huge function. It would be good to move its components outside. For example, get_decorator and visit_static_for should be (static) member functions, at the same level of visit_For itself.

This means you might have to pass a few originally automatically captured variables (e.g. node) as member function parameters, but doing so will further clarify this function. Thanks!

examples/mgpcg_advanced.py

yuanming-hu

Awesome! This is really well done. I like your TaichiSyntaxError treatments.

I'll merge once CI passes. Let me know when you are ready to chat about the next task.

yuanming-hu and others added 8 commits March 23, 2020 15:30

New extension 'adstack'

8db682f

[skip ci] enforce format

59354a4

New extension 'adstack' (taichi-dev#646)

ae3d22d

Simplify init in mgpcg_advanced.py (taichi-dev#645)

063542f

put static/range/struct for into separate functions

b077a71

make decorators look clearer

1110512

[skip ci] raise syntax error for more cases

75c90f6

[skip ci] dynamic grouped ndrange?

d8a3dae

xumingkuan commented Mar 24, 2020

View reviewed changes

yuanming-hu reviewed Mar 24, 2020

View reviewed changes

yuanming-hu force-pushed the master branch from 063542f to 460416e Compare March 24, 2020 02:32

xumingkuan added 3 commits March 24, 2020 14:40

[skip ci] debug

2a07300

fix

774ef9a

add more tests

72c4e4b

xumingkuan marked this pull request as ready for review March 24, 2020 18:58

xumingkuan and others added 4 commits March 24, 2020 15:05

Merge branch 'master' into for

ee8face

# Conflicts: # python/taichi/lang/transformer.py

[skip ci] enforce code format

02309ad

further simplify mgpcg_advanced.py and make it able to run with self.…

569b7ed

…dim=2

Merge branch 'for' of github.com:xumingkuan/taichi into for

d464c47

xumingkuan changed the title ~~Refactor ASTTransformer.visit_For~~ Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops Mar 24, 2020

[skip ci] enforce code format

306152b

yuanming-hu requested changes Mar 24, 2020

View reviewed changes

examples/mgpcg_advanced.py Outdated Show resolved Hide resolved

[skip ci] rename ijk to indices

224663b

move components of visit_For outside

25f3c7d

xumingkuan requested a review from yuanming-hu March 24, 2020 21:22

yuanming-hu approved these changes Mar 24, 2020

View reviewed changes

yuanming-hu merged commit 76ff830 into taichi-dev:master Mar 25, 2020

xumingkuan deleted the for branch March 26, 2020 03:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops #648

Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops #648

xumingkuan commented Mar 23, 2020 •

edited

Loading

xumingkuan commented Mar 24, 2020

xumingkuan Mar 24, 2020

yuanming-hu Mar 24, 2020

xumingkuan Mar 24, 2020

yuanming-hu Mar 24, 2020 •

edited

Loading

yuanming-hu Mar 24, 2020

yuanming-hu Mar 24, 2020

yuanming-hu commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

yuanming-hu commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

yuanming-hu left a comment

yuanming-hu left a comment

Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops #648

Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange loops #648

Conversation

xumingkuan commented Mar 23, 2020 • edited Loading

xumingkuan commented Mar 24, 2020

xumingkuan Mar 24, 2020

Choose a reason for hiding this comment

yuanming-hu Mar 24, 2020

Choose a reason for hiding this comment

xumingkuan Mar 24, 2020

Choose a reason for hiding this comment

yuanming-hu Mar 24, 2020 • edited Loading

Choose a reason for hiding this comment

yuanming-hu Mar 24, 2020

Choose a reason for hiding this comment

yuanming-hu Mar 24, 2020

Choose a reason for hiding this comment

yuanming-hu commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

yuanming-hu commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

xumingkuan commented Mar 24, 2020

yuanming-hu left a comment

Choose a reason for hiding this comment

yuanming-hu left a comment

Choose a reason for hiding this comment

xumingkuan commented Mar 23, 2020 •

edited

Loading

yuanming-hu Mar 24, 2020 •

edited

Loading